import os
import config.Yolov4 as cfg
import torch
from tqdm import tqdm

'''这个矩形训练暂时放弃，感觉不太好用'''

def __load_annotations(path):
    truth = {}
    if not os.path.isfile(cfg.dataset_truth):
        with open(path, 'r', encoding='utf-8') as f:
            pbar = f.readlines()
            pbar = tqdm(pbar)
            for line in pbar:
                data = line.split(' ')
                truth[data[0]] = []
                for i in data[1:]:
                    truth[data[0]].append([int(float(j)) for j in i.split(',')])
        torch.save(truth ,cfg.dataset_truth)
    else:
        truth = torch.load(cfg.dataset_truth)
    return truth

if __name__ == '__main__':
    cfg.dataset_truth = r'D:\python\upload\learn_yolov4_and_yolov5\logs\save.cache'
    __load_annotations(r'D:\python\upload\learn_yolov4_and_yolov5\data\\train_annotation.txt')